import matplotlib.pyplot as plt
import datetime
import tensorflow as tf
from tensorflow.keras import layers, backend, models, callbacks, datasets, preprocessing, optimizers, losses, metrics
from utils.time_print import printbar


class TrainMethods(object):
    """
    训练模型的三种方法
    """
    def __init__(self, batch_size=32, buffer_size=1000, max_len=300, embedding_output_dim=7):
        self.batch_size = batch_size
        self.buffer_size = buffer_size
        self.max_len = max_len
        self.embedding_output_dim = embedding_output_dim
        # 数据加载
        x_train, y_train, x_test, y_test = self.load_data()
        self.max_words = x_train.max() + 1
        self.cat_num = y_train.max() + 1
        # 数据管道
        self.ds_train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).\
            shuffle(buffer_size=buffer_size).\
            batch(batch_size=batch_size).\
            prefetch(tf.data.experimental.AUTOTUNE).cache()
        self.ds_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)). \
            shuffle(buffer_size=buffer_size). \
            batch(batch_size=batch_size). \
            prefetch(tf.data.experimental.AUTOTUNE).cache()

        self.model = self.create_model()

        self.loss_func = losses.SparseCategoricalCrossentropy()
        self.optimizer = optimizers.Nadam()
        self.train_metric = metrics.SparseCategoricalAccuracy(name='train_accuracy')
        self.test_metric = metrics.SparseCategoricalAccuracy(name='test_accuracy')

        self.train_loss = metrics.Mean(name='train_loss')
        self.test_loss = metrics.Mean(name='test_loss')

    def load_data(self):
        """
        数据加载
        :return:
        """
        (x_train, y_train), (x_test, y_test) = datasets.reuters.load_data()
        x_train = preprocessing.sequence.pad_sequences(sequences=x_train, maxlen=self.max_len)
        x_test = preprocessing.sequence.pad_sequences(sequences=x_test, maxlen=self.max_len)

        return x_train, y_train, x_test, y_test

    def create_model(self):
        """
        构建模型
        :return:
        """
        model = models.Sequential()

        model.add(layers.Embedding(input_dim=self.max_words, output_dim=self.embedding_output_dim,
                                   input_length=self.max_len))
        model.add(layers.Conv1D(filters=64, kernel_size=5, activation='relu'))
        model.add(layers.MaxPool1D(2))
        model.add(layers.Conv1D(filters=32, kernel_size=3, activation='relu'))
        model.add(layers.MaxPool1D(2))
        model.add(layers.Flatten())
        model.add(layers.Dense(units=self.cat_num, activation='softmax'))
        return model

    def compile_model(self):
        """
        编译model
        :param model:
        :return:
        """
        self.model.compile(optimizer=self.optimizer, loss=self.loss_func,
                      metrics=[self.train_metric, metrics.SparseTopKCategoricalAccuracy(5)])

    def train_model_on_batch(self, epochs):
        """
        方法二 train_on_batch 的train_model
        :param epochs:
        :return:
        """
        for epoch in tf.range(1, epochs + 1):
            self.model.reset_metrics()
            #  学习率调整
            if epoch == 5:
                self.model.optimizer.lr.assign(self.model.optimizer.lr / 2.0)
                tf.print('lowering optimizer learning rate')

            for x, y in self.ds_train:
                train_result = self.model.train_on_batch(x, y)

            for x, y in self.ds_test:
                test_result = self.model.test_on_batch(x, y)

            if epoch % 1 == 0:
                printbar()
                print('epoch = {}'.format(epoch))
                print('train: ', dict(zip(self.model.metrics_names, train_result)))
                print('test: ', dict(zip(self.model.metrics_names, test_result)))

    @tf.function
    def train_customize_step(self, features, labels):
        """
        方法三 自定义训练循环  train_step
        :param features:
        :param labels:
        :return:
        """
        with tf.GradientTape() as tape:
            predict = self.model(features, training=True)
            loss = self.loss_func(y_true=labels, y_pred=predict)
        gradients = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))

        self.train_loss.update_state(loss)
        self.train_metric.update_state(y_true=labels, y_pred=predict)

    @tf.function
    def test_customize_step(self, features, labels):
        """
        方法三 自定义训练循环  test_step
        :param features:
        :param labels:
        :return:
        """
        predict = self.model(features)
        batch_loss = self.loss_func(y_true=labels, y_pred=predict)
        self.test_loss.update_state(batch_loss)
        self.test_metric.update_state(y_true=labels, y_pred=predict)

    def train_model_by_customize(self, epochs):
        """
        方法三 自定义训练循环
        :param epochs:
        :return:
        """
        for epoch in tf.range(1, epochs + 1):
            for features, labels in self.ds_train:
                self.train_customize_step(features=features, labels=labels)

            for features, labels in self.ds_test:
                self.test_customize_step(features=features, labels=labels)

            logs = 'epoch={}, loss: {}, Accuracy: {}, test loss: {}, test accuracy: {}'

            if epoch % 1 == 0:
                printbar()
                tf.print(tf.strings.format(logs, (epoch, self.train_loss.result(), self.train_metric.result(),
                                                  self.test_loss.result(), self.test_metric.result())))

            self.train_loss.reset_states()
            self.test_loss.reset_states()
            self.train_metric.reset_states()
            self.test_metric.reset_states()

    def train_model_fit(self):
        """
        方法一：内置fit
        :return:
        """
        # 编译model
        self.compile_model()
        self.model.fit(self.ds_train, validation_data=self.ds_test, epochs=10)

    def train_model_train_on_batch(self):
        """
        方法二：内置train_on_batch
        :return:
        """
        # 编译model
        self.compile_model()
        self.train_model_on_batch(epochs=10)

    def train_model_customize(self):
        """
        方法三 自定义训练循环
        :return:
        """
        self.train_model_by_customize(epochs=10)


def run():
    train_model = TrainMethods()

    # 方法一：内置fit
    # train_model.train_model_fit()

    # 方法二：内置train_on_batch
    # train_model.train_model_train_on_batch()

    # 方法三 自定义训练循环
    train_model.train_model_customize()


if __name__ == '__main__':
    run()
